package com.bootx.controller;

import com.bootx.MyDocumentSplit;
import com.bootx.agent.TranslationAgent;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentParser;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.StreamingResponseHandler;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.injector.DefaultContentInjector;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.List;

/**
 * @author black
 */
@RestController
@RequestMapping("/ai")
public class AiController {

    @Resource(name = "deepSeekStreamingChatModel")
    private StreamingChatLanguageModel streamingChatModel1;

    @Resource(name = "ollamaEmbeddingModel")
    private EmbeddingModel embeddingModel;

    @Resource(name = "elasticsearchEmbeddingStore")
    private EmbeddingStore<TextSegment> embeddingStore;

    private final ChatMemory chatMemory = MessageWindowChatMemory.builder().maxMessages(20).build();


    @PostMapping(value = "/chat",produces = MediaType.TEXT_EVENT_STREAM_VALUE+";charset=UTF-8")
    public Flux<String> chat(String content, HttpServletResponse response,String[] filePaths) {
        response.setContentType(MediaType.TEXT_EVENT_STREAM_VALUE+";charset=UTF-8");
        UserMessage userMessage = UserMessage.from(content);
        if(filePaths.length>0) {
            for (String s : filePaths) {
                //1. 解析
                DocumentParser parser = new TextDocumentParser();
                Document document = FileSystemDocumentLoader.loadDocument("d:\\"+s, parser);
                MyDocumentSplit split = new MyDocumentSplit();
                List<TextSegment> textSegments = split.split(document);
                //2. 转换
                Response<List<Embedding>> listResponse = embeddingModel.embedAll(textSegments);
                //3. 存储
                embeddingStore.addAll(listResponse.content(), textSegments);
            }
            //4. 查询
            ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                    .embeddingStore(embeddingStore)
                    .minScore(0.8).
                    maxResults(10)
                    .embeddingModel(embeddingModel)
                    .build();
            List<Content> retrieve = contentRetriever.retrieve(Query.from(userMessage.singleText()));
            DefaultContentInjector defaultContentInjector = new DefaultContentInjector();
            UserMessage userMessage1 = defaultContentInjector.inject(retrieve, userMessage);
            chatMemory.add(userMessage1);
        }else {
            chatMemory.add(userMessage);
        }

        // 使用 Flux.create 来包装异步回调数据
        return Flux.create(emitter -> {
            // 使用流式的 chat 方法
            streamingChatModel1.chat(chatMemory.messages(), new StreamingChatResponseHandler() {

                @Override
                public void onPartialResponse(String partialResponse) {
                    emitter.next(partialResponse);
                }

                @Override
                public void onCompleteResponse(ChatResponse completeResponse) {
                    emitter.complete();
                }

                @Override
                public void onError(Throwable throwable) {
                    // 错误时发出 error 信号
                    emitter.error(throwable);
                }
            });
        });
    }



    @PostMapping(value = "/agent",produces = MediaType.TEXT_EVENT_STREAM_VALUE+";charset=UTF-8")
    public Flux<String> translation(Long agentId,String content, HttpServletResponse response,String language) {
        if(agentId==1){
            TranslationAgent assistant = AiServices.create(TranslationAgent.class, streamingChatModel1);
            return assistant.run(content, language);
        } else if (agentId==2) {
            TranslationAgent assistant = AiServices.create(TranslationAgent.class, streamingChatModel1);
            return assistant.run1(content);
        }
        TranslationAgent assistant = AiServices.create(TranslationAgent.class, streamingChatModel1);
        return assistant.run2(content);
    }
}
